#Copyright (c) 2025, Alibaba Cloud and its affiliates;
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.



from flask import Blueprint
from flask import current_app as app

from webargs import fields
from marshmallow import INCLUDE
from webargs.flaskparser import use_args
from marshmallow import Schema, fields

from services.cluster import clusterService
from common_libs.helper import response_data

cluster_controller = Blueprint('cluster', __name__, url_prefix='/cluster')

@cluster_controller.route('/envInit', methods=['POST'])
@use_args({
    "RequestId": fields.Str(required=True),
    "SDKPath": fields.Str(required=True),
    "KubeConfPath": fields.Str(required=False),
    "OperatorImage": fields.Str(required=True),
    "Namespace": fields.Str(required=False),
}, location="form", unknown=INCLUDE)
#@authorize() TBD
def envInit(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "SDKPath": args.get('SDKPath', ''),
        "KubeConfPath": args.get('KubeConfPath', '/root/.kube/config'),
        "OperatorImage": args.get('OperatorImage', ''),
        "Namespace": args.get('Namespace', 'default'),
    }

    app.logger.info('RequestId:{}, list ray clusters args: {}'.format(args['RequestId'], req))
    return clusterService.init(req)

@cluster_controller.route('/listRayClusters', methods=['POST'])
@use_args({
    "RequestId": fields.Str(required=True),
    "ClusterNames": fields.List(fields.Str(required=False)),
    "ClusterId": fields.Str(required=False),
    "ClusterVersion": fields.Str(required=False),
    "Namespace": fields.Str(required=False),
    "KubeConfPath": fields.Str(required=False),
}, location="form", unknown=INCLUDE)
#@authorize() TBD
def listRayClusters(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "ClusterNames": args.get('ClusterNames', []),
        "ClusterId": args.get('ClusterId', ''),
        "ClusterVersion": args.get('ClusterVersion', ''),
        "Namespace": args.get('Namespace', 'default'),
        "KubeConfPath": args.get('KubeConfPath', '/root/.kube/config')
    }

    app.logger.info('RequestId:{}, list ray clusters args: {}'.format(args['RequestId'], req))
    return clusterService.list_ray_cluster(req)

class EnvironmentVarsSchema(Schema):
    Name = fields.Str(required=True)
    Value = fields.Str(required=True)

class NodeAffinitySchema(Schema):
    Key = fields.Str(required=True)
    Operator = fields.Str(required=True)
    Values = fields.List(fields.Str(required=True), required=True)

class WorkerNodeGroupsSchema(Schema):
    Image = fields.Str(required=True)
    Cpu = fields.Int(required=False, missing=2)
    Memory = fields.Str(required=False, missing='4G')
    Replicas = fields.Int(required=False, missing=1)
    GroupName = fields.Str(required=False, missing='workergroup')
    NodeAffinity = fields.List(fields.Nested(NodeAffinitySchema), required=False, missing=[])
    EnvironmentVars = fields.List(fields.Nested(EnvironmentVarsSchema), required=False, missing=[])

class HeadNodeSchema(Schema):
    Image = fields.Str(required=True)
    Cpu = fields.Int(required=False, missing=1)
    Memory = fields.Str(required=False, missing='2G')
    NodeAffinity = fields.List(fields.Nested(NodeAffinitySchema), required=False, missing=[])
    EnvironmentVars = fields.List(fields.Nested(EnvironmentVarsSchema), required=False, missing=[])
    
class RequestSchema(Schema):
    RequestId = fields.Str(required=True)
    ClusterName = fields.Str(required=True)
    ClusterVersion = fields.Str(required=False)
    Description = fields.Str(required=False)
    KubeConfPath = fields.Str(required=False)
    Namespace = fields.Str(required=False)
    #OperatorImage = fields.Str(required=True)
    HeadNode = fields.Nested(HeadNodeSchema, required=True)
    WorkerNodeGroups = fields.List(fields.Nested(WorkerNodeGroupsSchema, required=True))

@cluster_controller.route('/createRayCluster', methods=['POST'])
@use_args(RequestSchema(), location="json", unknown=INCLUDE)
def createRayCluster(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "ClusterName": args.get('ClusterName', ''),
        "ClusterVersion": args.get('ClusterVersion', ''),
        "Description": args.get('Description', ''),
        "KubeConfPath": args.get('KubeConfPath', '/root/.kube/config'),
        "Namespace": args.get('Namespace', 'default'),
        "HeadNode": args.get('HeadNode', {}),
        "WorkerNodeGroups": args.get('WorkerNodeGroups', []),
    }

    app.logger.info('RequestId:{}, create ray clusters args: {}'.format(args['RequestId'], req))
    return clusterService.create_ray_cluster(req)


@cluster_controller.route('/deleteRayCluster', methods=['POST'])
@use_args({
    "RequestId": fields.Str(required=True),
    "KubeConfPath": fields.Str(required=False),
    "ClusterId": fields.Str(required=True),
    #"ClusterName": fields.Str(required=False),
    "Namespace": fields.Str(required=False),
}, location="form", unknown=INCLUDE)
#@authorize() TBD
def deleteRayCluster(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "KubeConfPath": args.get('KubeConfPath','/root/.kube/config'),
        "ClusterId": args.get('ClusterId',''),
        #"ClusterName": args.get('ClusterName',''),
        "Namespace": args.get('Namespace','default'),
    }

    app.logger.info('RequestId:{}, delete ray clusters args: {}'.format(args['RequestId'], req))
    return clusterService.delete_ray_cluster(req)

@cluster_controller.route('/getRayCluster', methods=['POST'])
@use_args({
    "RequestId": fields.Str(required=True),
    "KubeConfPath": fields.Str(required=False),
    "ClusterId": fields.Str(required=True),
}, location="form", unknown=INCLUDE)
def getRayCluster(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "KubeConfPath": args.get('KubeConfPath', '/root/.kube/config'),
        "ClusterId": args.get('ClusterId', ''),
    }

    app.logger.info('RequestId:{}, get ray cluster args: {}'.format(args['RequestId'], req))
    return clusterService.get_ray_cluster(req)

# 测试接口
@cluster_controller.route('/test', methods=['POST'])
@use_args({
    "RequestId": fields.Str(required=True),
    "Test": fields.Str(required=False)
}, location="form", unknown=INCLUDE)
def test(args):
    req = {
        "RequestId": args.get('RequestId',''),
        "Test": args.get('Test',''),
    }

    app.logger.info('RequestId:{}, test args: {}'.format(args['RequestId'], req))
    
    return response_data(True, req, "Success")